import matplotlib

matplotlib.use('TkAgg')  # 设置matplotlib后端为TkAgg，解决可能的显示问题
import math
import random
import matplotlib.pyplot as plt
import numpy as np

# 设置中文字体为黑体，解决中文显示可能出现的乱码问题，同时解决负号显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# 加载Iris数据集，返回特征和标签的元组列表
def load_iris_data():
    iris_data = []
    try:
        with open('iris.data', 'r') as file:
            for line in file:
                if line.strip():
                    row = line.strip().split(',')
                    features = list(map(float, row[:-1]))
                    label = row[-1]
                    iris_data.append((features, label))
    except FileNotFoundError:
        print("数据集文件不存在，请检查文件路径是否正确。")
    return iris_data


# 计算欧氏距离
def euclidean_distance(point1, point2):
    """
    计算两个数据点之间的欧氏距离。

    参数:
    point1 (list): 第一个数据点的特征向量
    point2 (list): 第二个数据点的特征向量

    返回:
    float: 两个数据点之间的欧氏距离
    """
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(point1, point2)))


# KNN分类算法，根据训练集对给定测试点进行分类
def knn_classify(train_data, test_point, k):
    """
    使用KNN算法对测试点进行分类。

    参数:
    train_data (list): 训练数据集，包含特征和标签的元组列表
    test_point (list): 待分类的测试点特征向量
    k (int): 选择的最近邻数量

    返回:
    str: 分类预测的标签
    """
    distances = []
    for features, label in train_data:
        distance = euclidean_distance(features, test_point)
        distances.append((distance, label))
    distances.sort(key=lambda x: x[0])  # 按距离排序
    nearest_neighbors = distances[:k]
    labels = [label for _, label in nearest_neighbors]
    return max(set(labels), key=labels.count)  # 多数表决


# 划分训练集和测试集
def train_test_split(iris_data, test_size=0.2):
    """
    将数据集划分为训练集和测试集。

    参数:
    iris_data (list): 完整的Iris数据集
    test_size (float): 测试集所占比例，默认0.2

    返回:
    tuple: 划分后的训练集和测试集
    """
    random.shuffle(iris_data)
    split_index = int(len(iris_data) * (1 - test_size))
    train_set = iris_data[:split_index]
    test_set = iris_data[split_index:]
    return train_set, test_set


# 主函数
if __name__ == "__main__":
    # 加载数据
    iris_data = load_iris_data()
    train_data, test_data = train_test_split(iris_data)

    k = 3  # 最近邻数量
    correct = 0
    total = len(test_data)

    # 用于存储不同类别数据点的坐标，方便绘图
    setosa_x = []
    setosa_y = []
    versicolor_x = []
    versicolor_y = []
    virginica_x = []
    virginica_y = []

    for features, label in test_data:
        prediction = knn_classify(train_data, features, k)
        if prediction == label:
            correct += 1

        # 根据真实类别将坐标分别存储到对应列表，用于绘图
        if label == "Iris-setosa":
            setosa_x.append(features[0])
            setosa_y.append(features[1])
        elif label == "Iris-versicolor":
            versicolor_x.append(features[0])
            versicolor_y.append(features[1])
        else:
            virginica_x.append(features[0])
            virginica_y.append(features[1])

    accuracy = correct / total
    print(f"KNN分类准确率: {accuracy * 100:.2f}%")

    # 绘制分类图
    plt.scatter(setosa_x, setosa_y, c='r', label='山鸢尾（Iris-setosa）')
    plt.scatter(versicolor_x, versicolor_y, c='g', label='变色鸢尾（Iris-versicolor）')
    plt.scatter(virginica_x, virginica_y, c='b', label='维吉尼亚鸢尾（Iris-virginica）')
    plt.xlabel('特征1')
    plt.ylabel('特征2')
    plt.title('Iris数据集KNN分类结果可视化')
    plt.legend()
    plt.show()
